from blocks.mobilenetv3 import PLA_MobileNetV3_Small
from blocks.NAFNet_arch import NAFNet
from blocks.unet import UNet
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
#import pytorch_colors as colors
import numpy as np


class enhance_net_nopool(nn.Module):

    def __init__(self):
        super(enhance_net_nopool, self).__init__()
        self.relu = nn.ReLU(inplace=True)
        number_f = 32
        self.e_conv1 = nn.Conv2d(3, number_f, 3, 1, 1, bias=True)
        self.e_conv2 = nn.Conv2d(number_f, number_f, 3, 1, 1, bias=True)
        self.e_conv3 = nn.Conv2d(number_f, number_f, 3, 1, 1, bias=True)
        self.e_conv4 = nn.Conv2d(number_f, number_f, 3, 1, 1, bias=True)
        self.e_conv5 = nn.Conv2d(number_f*2, number_f, 3, 1, 1, bias=True)
        self.e_conv6 = nn.Conv2d(number_f*2, number_f, 3, 1, 1, bias=True)
        self.e_conv7 = nn.Conv2d(number_f*2, 24, 3, 1, 1, bias=True)

        self.maxpool = nn.MaxPool2d(2, stride=2, return_indices=False, ceil_mode=False)
        self.upsample = nn.UpsamplingBilinear2d(scale_factor=2)
        self.backbone = UNet()

    def forward(self, x):

        # x1 = self.relu(self.e_conv1(x))
        # # p1 = self.maxpool(x1)
        # x2 = self.relu(self.e_conv2(x1))
        # # p2 = self.maxpool(x2)
        # x3 = self.relu(self.e_conv3(x2))
        # # p3 = self.maxpool(x3)
        # x4 = self.relu(self.e_conv4(x3))

        # x5 = self.relu(self.e_conv5(torch.cat([x3, x4], 1)))
        # # x5 = self.upsample(x5)
        # x6 = self.relu(self.e_conv6(torch.cat([x2, x5], 1)))
        # x7 = self.e_conv7(torch.cat([x1, x6], 1))
        # # x7 = self.pla_cbam(x7)
        # x_r = F.tanh(x7)
        x_r = self.backbone(x)
        r1, r2, r3, r4, r5, r6, r7, r8 = torch.split(x_r, 3, dim=1)

        #
        x = x + r1*(torch.pow(x, 2)-x)
        x = x + r2*(torch.pow(x, 2)-x)
        x = x + r3*(torch.pow(x, 2)-x)
        enhance_image_1 = x + r4*(torch.pow(x, 2)-x)
        x = enhance_image_1 + r5*(torch.pow(enhance_image_1, 2)-enhance_image_1)
        x = x + r6*(torch.pow(x, 2)-x)
        x = x + r7*(torch.pow(x, 2)-x)
        enhance_image = x + r8*(torch.pow(x, 2)-x)
        r = torch.cat([r1, r2, r3, r4, r5, r6, r7, r8], 1)
        return enhance_image_1, enhance_image, r


class DA_DCE_Net(nn.Module):
    def __init__(self):
        super(DA_DCE_Net, self).__init__()
        self.zero_dce = enhance_net_nopool()
        # enc_blks = [2, 2, 4, 8]
        # middle_blk_num = 12
        # dec_blks = [2, 2, 2, 2]

        enc_blks = [1, 1, 2, 8]
        middle_blk_num = 6
        dec_blks = [1, 1, 1, 1]
        self.denoizing = NAFNet(middle_blk_num=middle_blk_num, enc_blk_nums=enc_blks, dec_blk_nums=dec_blks)

    def forward(self, x):
        enhance_image_1, enhance_image, r = self.zero_dce(x)
        denoise_result = self.denoizing(enhance_image*255)
        return enhance_image_1, enhance_image, denoise_result, r
